import numpy as np
import struct
import os
import scipy.io
import time
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from layers_1 import FullyConnectedLayer, ReLULayer, SoftmaxLossLayer
from layers_2 import ConvolutionalLayer, MaxPoolingLayer, FlattenLayer

def show_matrix(mat, name):
    #print(name + str(mat.shape) + ' mean %f, std %f' % (mat.mean(), mat.std()))
    pass

class VGG19(object):
    def __init__(self, param_path='../imagenet-vgg-verydeep-19.mat'):
        # 神经网络的初始化
        self.param_path = param_path
        self.param_layer_name = (
            'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',
            'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
            'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3',
            'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4',
            'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5',
            'flatten', 'fc6', 'relu6', 'fc7', 'relu7', 'fc8', 'Softmax'
        )
        self.layers = {}
        self.update_layer_list = []

    def build_model(self):
        # 定义VGG19的网络结构
        self.layers['conv1_1'] = ConvolutionalLayer(3, 3, 64, 1, 1)  # [in_ch, out_ch, kernel, pad, stride]
        self.layers['relu1_1'] = ReLULayer()
        self.layers['conv1_2'] = ConvolutionalLayer(3, 64, 64, 1, 1)
        self.layers['relu1_2'] = ReLULayer()
        self.layers['pool1'] = MaxPoolingLayer(2, 2)  # kernel=2, stride=2
        # TODO 中间层省略（需按param_layer_name补全）
        # Block 2
        self.layers['conv2_1'] = ConvolutionalLayer(3, 64, 128, 1, 1)
        self.layers['relu2_1'] = ReLULayer()
        self.layers['conv2_2'] = ConvolutionalLayer(3, 128, 128, 1, 1)
        self.layers['relu2_2'] = ReLULayer()
        self.layers['pool2'] = MaxPoolingLayer(2, 2)
        # Block 3
        self.layers['conv3_1'] = ConvolutionalLayer(3, 128, 256, 1, 1)
        self.layers['relu3_1'] = ReLULayer()
        self.layers['conv3_2'] = ConvolutionalLayer(3, 256, 256, 1, 1)
        self.layers['relu3_2'] = ReLULayer()
        self.layers['conv3_3'] = ConvolutionalLayer(3, 256, 256, 1, 1)
        self.layers['relu3_3'] = ReLULayer()
        self.layers['conv3_4'] = ConvolutionalLayer(3, 256, 256, 1, 1)
        self.layers['relu3_4'] = ReLULayer()
        self.layers['pool3'] = MaxPoolingLayer(2, 2)
        # Block 4
        self.layers['conv4_1'] = ConvolutionalLayer(3, 256, 512, 1, 1)
        self.layers['relu4_1'] = ReLULayer()
        self.layers['conv4_2'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu4_2'] = ReLULayer()
        self.layers['conv4_3'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu4_3'] = ReLULayer()
        self.layers['conv4_4'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu4_4'] = ReLULayer()
        self.layers['pool4'] = MaxPoolingLayer(2, 2)
        # Block 5
        self.layers['conv5_1'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu5_1'] = ReLULayer()
        self.layers['conv5_2'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu5_2'] = ReLULayer()
        self.layers['conv5_3'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu5_3'] = ReLULayer()
        # ______________________
        self.layers['conv5_4'] = ConvolutionalLayer(3, 512, 512, 1, 1)
        self.layers['relu5_4'] = ReLULayer()
        self.layers['pool5'] = MaxPoolingLayer(2, 2)
        # Block 6
        self.layers['flatten'] = FlattenLayer([512, 7, 7], [512*7*7])  # 输入形状需根据实际计算调整
        self.layers['fc6'] = FullyConnectedLayer(512*7*7, 4096)
        self.layers['relu6'] = ReLULayer()
        # TODO
        self.layers['fc7'] = FullyConnectedLayer(4096, 4096)  # 原代码缺失，此处补充
        self.layers['relu7'] = ReLULayer()
        # ______________________
        self.layers['fc8'] = FullyConnectedLayer(4096, 1000)  # ImageNet分类数
        self.layers['Softmax'] = SoftmaxLossLayer()

        # 记录需要更新参数的层（卷积层和全连接层）
        self.update_layer_list = [
            layer_name for layer_name in self.layers.keys() 
            if 'conv' in layer_name or 'fc' in layer_name
        ]
    
    def init_model(self):
        print('Initializing parameters of each layer in vgg-19...')
        for layer_name in self.update_layer_list:
            self.layers[layer_name].init_param()

    def load_model(self):  # 加载预训练参数
        params = scipy.io.loadmat(self.param_path)
        self.image_mean = params['normalization'][0][0][0]
        self.image_mean = np.mean(self.image_mean, axis=(0, 1))  # 计算图像均值 [C]

        for idx in range(43):  # 遍历所有层（VGG19共43层）
            layer_name = self.param_layer_name[idx]
            
            # 加载卷积层参数
            if 'conv' in layer_name:
                weight, bias = params['layers'][0][idx][0][0][0][0]
                # 调整权重维度顺序：从[H,W,in_C,out_C] -> [in_C,H,W,out_C]
                weight = np.transpose(weight, [2, 0, 1, 3])
                bias = bias.reshape(-1)  # 展平偏置
                self.layers[layer_name].load_param(weight, bias)
            
            # 加载全连接层参数（fc6,fc7,fc8）
            elif idx >= 37 and 'fc' in layer_name:  # 最后6层中的全连接层
                weight, bias = params['layers'][0][idx-1][0][0][0][0]
                # 调整全连接层权重：[H,W,in_C,out_C] -> [in_C*H*W, out_C]
                weight = weight.reshape([weight.shape[0]*weight.shape[1]*weight.shape[2], weight.shape[3]])
                self.layers[layer_name].load_param(weight, bias)

    def load_image(self, image_dir):
        print('Loading and preprocessing image from ' + image_dir)
        self.input_image = scipy.misc.imread(image_dir)
        self.input_image = scipy.misc.imresize(self.input_image,[224,224,3])
        self.input_image = np.array(self.input_image).astype(np.float32)
        self.input_image -= self.image_mean
        # input dim [N, channel, height, width]
        # TODO：调整图片维度顺序
        self.input_image = np.transpose(self.input_image, [2, 0, 1]) 
        # ________________________
        self.input_image = np.reshape(self.input_image, [1]+list(self.input_image.shape))

    def forward(self, input_image):  # 前向传播
        self.input_image = input_image - self.image_mean  # 减去均值
        current = self.input_image
        for layer_name in self.param_layer_name:  # 按顺序通过各层
            current = self.layers[layer_name].forward(current)
        return current  # 输出Softmax概率 [N, 1000]

    def evaluate(self, input_image):  # 分类推断
        prob = self.forward(input_image)
        top1 = np.argmax(prob[0])  # 取概率最高的类别
        print('Classification result: id = %d, prob = %f' % (top1, prob[0, top1]))
        return top1

if __name__ == '__main__':
    vgg = VGG19()
    vgg.build_model()
    vgg.init_model()
    vgg.load_model()
    vgg.load_image('../../cat1.jpg')
    prob = vgg.evaluate()
    
